﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Text;
using msilgc.cbuilder.CSyntaxTree;
using msilgc.common;
using msilgc.driver.cuda.Compilers;

namespace msilgc.driver.cuda {
    internal class CudaCodeGenerator {

        private static readonly Dictionary<string, string> _structTypes =
            new Dictionary<string, string>
                {
                    {"int", "int_array"},
                    {"unsigned int", "unsigned_int_array"},
                    {"short", "short_array"},
                    {"unsigned short", "unsigned_short_array"},
                    {"long int", "long_int_array"},
                    {"unsigned long int", "unsigned_long_int_array"},
                    {"float", "float_array"},
                    {"double", "double_array"},
                    {"char", "char_array"}
                };

        public void GenerateNativeLibrary(DirectoryInfo dllOutputPath, IEnumerable<ICMethodBody> methodBody, out string compilerMsgs,
                                          out string kernelCode, bool writeDebugCode) {
            Logger.WriteLine("\tGenerating Code...");
            var kernels = WriteCode(methodBody, writeDebugCode);
            kernelCode = string.Join(Environment.NewLine, kernels);
            EmitHelperAndIncludeFiles(dllOutputPath);

            var files = new List<FileInfo>();
            var index = (-1);
            foreach (var kernel in kernels) {
                var file = new FileInfo(Path.Combine(dllOutputPath.FullName, "driver" + ++index + ".cu"));
                files.Add(file);
                if (File.Exists(file.FullName)) {
                    File.Delete(file.FullName);
                }
                File.WriteAllText(file.FullName, kernel);
            }

            string compilerMsg1;
            CompilerFactory.Create().RunCompiler(dllOutputPath, files, out compilerMsg1);
            compilerMsgs = compilerMsg1;
			CompilerFactory.Create().RunLinker(dllOutputPath, files.Select(x => x.Name.Replace(".cu", "")).ToArray());
            
        }

        public void GenerateNativeLibrary(DirectoryInfo dllOutputPath, ICMethodBody methodBody, out string compilerMsgs, out string kernelCode, bool writeDebugCode) {
            GenerateNativeLibrary(dllOutputPath, new[]{methodBody}, out compilerMsgs, out kernelCode, writeDebugCode);
        }
		
		private static void EmitHelperAndIncludeFiles(DirectoryInfo dllOutputPath) {
			File.WriteAllBytes(Path.Combine(dllOutputPath.FullName, "cuPrintf.cu"), Resource1.cuPrintf_cu);
			File.WriteAllBytes(Path.Combine(dllOutputPath.FullName, "cuPrintf.cuh"), Resource1.cuPrintf_cuh);
		}

        private static IEnumerable<IVariable> WriteCopyArguments(ICMethodBody methodBody, StringBuilder code, bool writeDebugCode) {
            var deviceArguments = new List<IVariable>();
            foreach (IVariable argument in methodBody.Arguments) {
                if (argument == null) {
                    continue;
                }
                if (!argument.IsPrimitive) {
                    if (writeDebugCode) {
                        code.Append("printf(\"allocating device array of size = %d\\r\\n\", ");
                        code.Append(argument.Name);
                        code.AppendLine(".array_length);");
                    }
                    if (writeDebugCode) {
                        code.Append("_debug_print(");
                        code.Append(argument.Name);
                        code.AppendLine(");");
                    }
                    IVariable deviceVariable = argument.CloneWithPrefix("d_");
                    deviceVariable.EmitDeclaration(code);
                    code.AppendLine(";");
                    code.Append("cutilSafeCall( cudaMalloc( (void**) &");
                    code.Append(deviceVariable.Name);
                    code.Append(".array_data");
                    code.Append(", sizeof(");
                    code.Append(deviceVariable.ArrayBaseType);
                    code.Append(") * ");
                    code.Append(argument.Name);
                    code.Append(".array_length");
                    code.AppendLine("));");
                    code.Append("cutilSafeCall( cudaMemcpy( ");
                    code.Append(deviceVariable.Name);
                    code.Append(".array_data, ");
                    code.Append(argument.Name);
                    code.Append(".array_data, ");
                    code.Append(argument.Name);
                    code.Append(".array_length * sizeof(");
                    code.Append(deviceVariable.ArrayBaseType);
                    code.AppendLine("), cudaMemcpyHostToDevice));");
                    code.AppendLine();
                    code.Append(deviceVariable.Name);
                    code.Append(".array_length = ");
                    code.Append(argument.Name);
                    code.Append(".array_length;");
                    code.AppendLine();
                    deviceVariable.OutVariable = argument.OutVariable;
                    deviceArguments.Add(deviceVariable);
                } else {
                    deviceArguments.Add(argument);
                }
            }
            return deviceArguments;
        }

        private static void WriteArguments(IEnumerable<IVariable> arguments, StringBuilder code) {
            bool first = true;
            foreach (var argument in arguments) {
                if (first) {
                    first = false;
                } else {
                    code.Append(", ");
                }
                code.Append(argument.Name);
            }
        }

        private static void EmitArrayTypes(StringBuilder code) {
            foreach (string typeKey in _structTypes.Keys) {
                code.AppendLine("typedef struct {");
                code.AppendLine("int array_length;");
                code.Append(typeKey); code.AppendLine(" *array_data;");
                code.Append("} "); code.Append(_structTypes[typeKey]); code.AppendLine(";");
                code.AppendLine();
                code.Append("void _debug_print("); code.Append(_structTypes[typeKey]); code.AppendLine(" arr){");
                code.AppendLine("for(int i = 0; i < arr.array_length; i++) { printf(\"%d\\n\", arr.array_data[i]); }");
                code.AppendLine("}");
            }
            code.AppendLine("typedef struct{ int x; int y; int z; } grain_type; ");
        }

        private static void EmitMethodDriverDeclaration(ICMethodBody methodBody, StringBuilder code) {
            methodBody.EmitReturnType(code);
            code.Append(" ");
            code.Append(methodBody.MethodName);
            code.Append("_driver(");
            methodBody.EmitArguments(code);
            code.Append(", grain_type threadgrainarg, grain_type datagrainarg, int deviceId)");
        }

        private static void WriteIncludeDeclarations(StringBuilder code) {
            code.AppendLine("#include <stdio.h>");
        }

        private static void AddUtilityMethods(StringBuilder code) {
            const string utilityMethods =
                "#define cutilSafeCall(err)           __cudaSafeCall      (err, __FILE__, __LINE__)\r\n" +
                "#define cutilSafeThreadSync()        __cudaSafeThreadSync(__FILE__, __LINE__)\r\n" +
                "__inline void __cudaSafeCall( cudaError err, const char *file, const int line )\r\n" +
                "{\r\n" +
                    "if( cudaSuccess != err) {\r\n" +
                        "fprintf(stderr, \"cudaSafeCall() Runtime API error in file <%s>, line %i : %s.\\n\", file, line, cudaGetErrorString( err) );\r\n" +
                        "exit(-1);\r\n" +
                    "}\r\n" +
                "}\r\n" +
                "__inline void __cudaSafeThreadSync( const char *file, const int line )\r\n" +
                "{\r\n" +
                    "cudaError err = cudaThreadSynchronize();\r\n" +
                    "if ( cudaSuccess != err) {\r\n" +
                        "fprintf(stderr, \"cudaThreadSynchronize() Driver API error in file '%s' in line %i : %s.\\n\", file, line, cudaGetErrorString( err) );\r\n" +
                        "exit(-1);\r\n" +
                    "}\r\n" +
                "}\r\n";
            code.Append(utilityMethods);
        }

        public IEnumerable<string> WriteCode(IEnumerable<ICMethodBody> methodBodies, bool writeDebugCode) {
            foreach (var cMethodBody in methodBodies) {
                yield return WriteCode(cMethodBody, writeDebugCode);
            }
        }

        public string WriteCode(ICMethodBody methodBody, bool writeDebugCode) {
            StringBuilder code = new StringBuilder(64 * 1024);
            WriteIncludeDeclarations(code);
            code.AppendLine();
            AddUtilityMethods(code);
            code.AppendLine();
            EmitArrayTypes(code);
            code.AppendLine();
            if (IsUnix()) {
                code.Append("extern \"C\" ");
            } else {
                code.Append("extern \"C\" __declspec(dllexport) ");
            }
            EmitMethodDriverDeclaration(methodBody, code);
            code.AppendLine(";");
            code.AppendLine();
            code.Append("__global__ ");
            methodBody.Emit(code, new CudaMethodNameMapper());
            code.AppendLine();

            EmitMethodDriverDeclaration(methodBody, code);
            code.AppendLine("{");
            code.AppendLine("cudaDeviceProp prop;");
            code.AppendLine("cudaGetDeviceProperties(&prop, deviceId);");
            if (writeDebugCode) {
                code.AppendLine("printf(\"using device name = %s, id=%d\\r\\n\", &prop.name[0], deviceId);");
            }
            code.AppendLine("cudaSetDevice(deviceId);");
            code.AppendLine();

            IEnumerable<IVariable> deviceArguments = WriteCopyArguments(methodBody, code, writeDebugCode);

            code.AppendLine("dim3 datagrain(datagrainarg.x, datagrainarg.y, datagrainarg.z);");
            code.AppendLine("dim3 threadgrain(threadgrainarg.x, threadgrainarg.y, threadgrainarg.z);");
            code.AppendLine();

            if (writeDebugCode) {
                code.AppendLine("printf(\"threadgrain : {%d,%d,%d}\\n\", threadgrainarg.x, threadgrainarg.y, threadgrainarg.z);");
                code.AppendLine("printf(\"datagrain : {%d,%d,%d}\\n\", datagrainarg.x, datagrainarg.y, datagrainarg.z);");
            }

            code.Append(methodBody.MethodName);
            code.Append("<<<datagrain,threadgrain>>>(");
            WriteArguments(deviceArguments, code);
            code.AppendLine(");");
            code.AppendLine();
            code.AppendLine("cudaThreadSynchronize();");
            code.AppendLine();
            bool first = true;
            foreach (IVariable argument in deviceArguments) {
                if (argument.OutVariable) {
                    if (first) {
                        first = false;
                        if (writeDebugCode) {
                            code.AppendLine("printf(\"copying back output variables...\\r\\n\");");
                        }
                    }
                    if (argument.IsPrimitive) {
                        throw new NotSupportedException("i don't yet support primitive out variable types. sorry!");
                    }
                    code.Append("cutilSafeCall(cudaMemcpy(");
                    code.Append(argument.Name.Substring(2));
                    code.Append(".array_data, ");
                    code.Append(argument.Name);
                    code.Append(".array_data, ");
                    code.Append(argument.Name.Substring(2));
                    code.Append(".array_length * sizeof(");
                    code.Append(argument.ArrayBaseType);
                    code.AppendLine("), cudaMemcpyDeviceToHost));");
                    if (writeDebugCode) {
                        code.Append("_debug_print(");
                        code.Append(argument.Name.Substring(2));
                        code.AppendLine(");");
                    }
                }
            }
            if (!first) {
                if (writeDebugCode) {
                    code.AppendLine("printf(\"done copying back output variables.\\r\\n\");");
                }
            }
            first = true;
            foreach (IVariable argument in deviceArguments) {
                if (!argument.IsPrimitive) {
                    if (first) {
                        first = false;
                        if (writeDebugCode) {
                            code.AppendLine("printf(\"freeing device memory...\\r\\n\");");
                        }
                    }
                    code.Append("cutilSafeCall(cudaFree(");
                    code.Append(argument.Name);
                    code.Append(".array_data");
                    code.AppendLine("));");
                }
            }
            if (!first) {
                if (writeDebugCode) {
                    code.AppendLine("printf(\"done releasing device memory.\\r\\n\");");
                }
            }

            code.AppendLine();
            code.AppendLine("cudaThreadExit();");
            code.AppendLine();

            code.AppendLine("}");
            code.AppendLine();

            return code.ToString();
        }

        private static bool IsUnix() {
            switch (Environment.OSVersion.Platform) {
                case PlatformID.Win32S:
                case PlatformID.Win32Windows:
                case PlatformID.Win32NT:
                return false;
                case PlatformID.Unix:
                case PlatformID.MacOSX:
                return true;
                default:
                throw new SystemException("platform is not supported");
            }
        }
    }
}